#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include <openssl/pem.h>
#include <openssl/rand.h>
#include <openssl/err.h>

#include "trace.h"
#include "util.h"

#define ERR_BUF_SIZE            (300u)
#define MAX_SIG_LEN             (1000u)

static void my_DSA_sign(const unsigned char *hash, int dlen, unsigned char *sig,
        int *siglen, DSA *dsa);
static int dsa_sign_setup_with_msg(DSA *dsa, BIGNUM **kinvp,
        BIGNUM **rp, const unsigned char *sub_msg, int sub_msg_len);
static int dsa_sign_setup_with_leak(DSA *dsa, BIGNUM **kinvp, BIGNUM **rp );

int dsa_sign_file(const char *filename,
        const char *prv_key_pem_file, const char *passphrase,
        const char *sub_msg, int sub_msg_len,
        int leak){
    int ret_code = 0;

    if((NULL == filename) || (NULL == prv_key_pem_file)) {
        TRACE_ERROR("Invalid input parameters\n");
        return 1;
    }

    // Step1 - load private key from file
    SSLeay_add_all_algorithms();
    DSA *dsa = NULL;
    if(NULL == passphrase) {
        // no passphrase provided. Ask
        if(load_private_key_dsa(prv_key_pem_file , "", &dsa)){
            TRACE_ERROR("load_private_key_dsa with password callback fail\n");
            return 2;
        }
    }
    else{
        if(load_private_key_dsa(prv_key_pem_file, passphrase, &dsa)){
            TRACE_ERROR("load_private_key_dsa with password fail\n");
            return 2;
        }
    }

    // Step 2 - compute sha1 hash of the file
    unsigned char *hash = (unsigned char *)malloc(SHA_DIGEST_LENGTH);
    if(NULL == hash) {
        TRACE_ERROR("malloc error\n");
        return 3;
    }

    unsigned int hash_len;
    if(calc_sha1(filename, &hash, &hash_len)){
        TRACE_ERROR("calc_sha1 error\n");
        return 3;
    }

    TRACE_DEBUG("SHA1(%s)=", filename);
    for(unsigned int i = 0; i < hash_len; i++){
        TRACE_DEBUG("%02x", hash[i]);
    }

    // Step 3 - sign the hash
    // 3.a - precompute  inv(k) and r
    BIGNUM *kinvp = NULL;
    BIGNUM *rp = NULL;

    if(NULL != sub_msg) {
        // a subliminal message to embed was provided
        if(dsa_sign_setup_with_msg(dsa, &kinvp, &rp,
                (const unsigned char *)sub_msg, sub_msg_len)){
            TRACE_WARNING("dsa_sign_setup_with_msg failed\n");
        }
    }
    else if (1 == leak){
        if(dsa_sign_setup_with_leak(dsa,&kinvp, &rp)){
            TRACE_WARNING("dsa_sign_setup_with_leak failed\n");
        }
    }
    else {
        // PRNG must be seeded before the call to DSA_sign_setup()
        RAND_egd("/dev/random");
        DSA_sign_setup(dsa, NULL, &kinvp, &rp);
    }

    dsa->r = rp;
    dsa->kinv = kinvp;

    // 3.b - sign and save signature from internal format to ASN1.DER
    unsigned char *out_sig = (unsigned char *)malloc(DSA_size(dsa));
    int siglen;

    my_DSA_sign(hash, 20, out_sig, &siglen, dsa);

    // 3.c - export DER to file
    BIO *out;
    out = BIO_new(BIO_s_file());
    if(out == NULL) {
        TRACE_ERROR("BIO_new error\n");
        return 1;
    }

    int sig_filename_len = strlen(filename)+5;
    char *sig_file = (char *)malloc(sig_filename_len);

    strncpy(sig_file, filename, strlen(filename));
    strncpy(sig_file+strlen(filename), ".sig", 4);

    sig_file[sig_filename_len-1] = '\0';

    if(!BIO_write_filename(out, sig_file)) {
        TRACE_ERROR("BIO_write_filename error\n");
        return 2;
    }

    if(siglen != BIO_write(out, out_sig, siglen)){
        TRACE_WARNING("Incorrect number of bytes writen to DER format\n");
    }

    TRACE_THIS("Written signature to %s file\n", sig_file);

    BIO_free(out);

    EVP_cleanup();

    return ret_code;
}

static void my_DSA_sign(const unsigned char *hash, int dlen, unsigned char *sig,
        int *siglen, DSA *dsa){
    // sign with raw signature (r,s)
    DSA_SIG *dsa_sig = DSA_SIG_new();

    RAND_seed(hash, dlen);
    dsa_sig = DSA_do_sign(hash, SHA_DIGEST_LENGTH, dsa);

    if(!dsa_sig){
        char buf[ERR_BUF_SIZE];
        unsigned long errcode = ERR_get_error();

        ERR_load_crypto_strings();
        ERR_error_string(errcode, buf);
        TRACE_ERROR("\nDSA_do_sign error: %s\n", buf);
        ERR_free_strings();
        return;
    }

    TRACE_DEBUG("Signature:\n");
    char *r;
    r = BN_bn2dec(dsa_sig->r);
    TRACE_DEBUG("r:%s\n", r);
    OPENSSL_free(r);

    char *s;
    s = BN_bn2dec(dsa_sig->s);
    TRACE_DEBUG("s:%s\n", s);
    OPENSSL_free(s);

    *siglen = i2d_DSA_SIG(dsa_sig, &sig);
    TRACE_DEBUG("Written signature with len: %d\n", *siglen);

    DSA_SIG_free(dsa_sig);
}

static int dsa_sign_setup_with_msg(DSA *dsa, BIGNUM **kinvp,
        BIGNUM **rp, const unsigned char *sub_msg, int sub_msg_len){
    BN_CTX *ctx;

    BIGNUM k, *kinv = NULL, *r = NULL;

    int ret = 0;

    if (!dsa->p || !dsa->q || !dsa->g){
        TRACE_ERROR("Invalid DSA parameters\n");
        return 1;
    }

    BN_init(&k);

    if ((ctx=BN_CTX_new()) == NULL) goto err;

    if ((r=BN_new()) == NULL) goto err;

    /* Write subliminal msg into k*/
    BN_bin2bn(sub_msg, sub_msg_len, &k);

    /* Compute r = (g^k mod p) mod q */
    if(!BN_mod_exp(r, dsa->g, &k, dsa->p, ctx)) goto err;
    if (!BN_mod(r, r, dsa->q, ctx)) goto err;

    TRACE_DEBUG("Precomputed r:\n");
    char *rr;
    rr = BN_bn2dec(r);
    TRACE_DEBUG("r:%s\n", rr);
    OPENSSL_free(rr);

    /* Compute  inv(k) */
    if ((kinv=BN_mod_inverse(NULL, &k, dsa->q, ctx)) == NULL) goto err;

    if (*kinvp != NULL) BN_clear_free(*kinvp);

    *kinvp=kinv;
    kinv=NULL;

    if (*rp != NULL) BN_clear_free(*rp);
    *rp=r;
    r=NULL;

    goto ok;

    err:
    ret = 1;

    ok:
    if (kinv != NULL) BN_clear_free(kinv);
    if (r != NULL) BN_clear_free(r);
    if (ctx != NULL) BN_CTX_free(ctx);

    BN_clear_free(&k);

    return(ret);
}

static int dsa_sign_setup_with_leak(DSA *dsa, BIGNUM **kinvp, BIGNUM **rp ){
    BN_CTX *ctx;

    BIGNUM k, *kinv = NULL, *r = NULL;

    int ret = 0;

    char *rr = NULL;
    unsigned char block_num;

    char secret_primes[14][4] = {"37", "79", "131", "181", "239",
                                 "293", "47", "97", "149", "197",
                                 "257", "313", "379", "439"};
    int to_leak[14] = {};

    if (!dsa->p || !dsa->q || !dsa->g){
        TRACE_ERROR("Invalid DSA parameters\n");
        return 1;
    }

    BN_init(&k);

    if ((ctx=BN_CTX_new()) == NULL) goto err;

    if ((r=BN_new()) == NULL) goto err;

    RAND_egd("/dev/random");

    TRACE_DEBUG("Number of (significant) bits in private key: %d\n", BN_num_bits(dsa->priv_key));
    /* Bit numbering in private key:
     * 77 5A 2C B1 1A 00 5A 6C 1A 34 36 6E 84 72 DB 2F F8 DE 7D 0E
     * 01110111 ... 00001110
     * bit159 bit 158 ...... bit2 bit1 bti0
     */
    char *to;
    to = BN_bn2hex(dsa->priv_key);
    TRACE_DEBUG("Private key: 0x%s\n", to);
    OPENSSL_free(to);

    /* 10 bit block number to leak (0..15) */
    srand (time(NULL));
    block_num =rand() % 16;
    for(int i=0; i<10; i++){
        to_leak[i] = BN_is_bit_set(dsa->priv_key, block_num*10+i);
    }
    for(int i=0; i<4; i++){
        to_leak[10+i] = block_num & (1 << i);
    }

    bool found;

    do{
        found = true;

        /* Generate random k */
        BN_init(&k);
        do{
            if (!BN_rand_range(&k, dsa->q)) goto err;
        }
        while (BN_is_zero(&k));

        /* Compute r = (g^k mod p) mod q */
        if(!BN_mod_exp(r, dsa->g, &k, dsa->p, ctx)) goto err;
        if (!BN_mod(r, r, dsa->q, ctx)) goto err;

        rr = BN_bn2dec(r);
        TRACE_DEBUG("Testing r:%s\n", rr);


        /* check quadratic residues conditions */
        for(int i=0; i<14; i++){
            if(to_leak[i] == is_quad_res(rr, secret_primes[i])){
                found = false;
                break;
            }
        }
    }while(!found);

    OPENSSL_free(rr);


    /* Compute  inv(k) */
    if ((kinv=BN_mod_inverse(NULL, &k, dsa->q, ctx)) == NULL) goto err;

    if (*kinvp != NULL) BN_clear_free(*kinvp);

    *kinvp=kinv;
    kinv=NULL;

    if (*rp != NULL) BN_clear_free(*rp);
    *rp=r;
    r=NULL;

    goto ok;

    err:
    ret = 1;

    ok:
    if (kinv != NULL) BN_clear_free(kinv);
    if (r != NULL) BN_clear_free(r);
    if (ctx != NULL) BN_CTX_free(ctx);

    BN_clear_free(&k);

    return(ret);
}

int dsa_get_k_with_x(const char *filename,
        const char *sig_file,
        const char *prv_key_pem_file,
        const char *passphrase,
        unsigned char **msg ){
    BIGNUM *w = BN_new();
    BN_CTX *ctx = BN_CTX_new();
    BIGNUM *m = BN_new();
    BIGNUM *k = BN_new();

    /* load private key from file */
    SSLeay_add_all_algorithms();
    DSA *dsa = NULL;
    if(NULL == passphrase) {
        // no passphrase provided. Ask
        if(load_private_key_dsa(prv_key_pem_file , "", &dsa)){
            TRACE_ERROR("load_private_key_dsa with password callback fail\n");
            return 2;
        }
    }
    else{
        if(load_private_key_dsa(prv_key_pem_file, passphrase, &dsa)){
            TRACE_ERROR("load_private_key_dsa with password fail\n");
            return 2;
        }
    }

    /* compute sha1 hash of the file */
    unsigned char *hash = (unsigned char *)malloc(SHA_DIGEST_LENGTH);
    if(NULL == hash) {
        TRACE_ERROR("malloc error\n");
        return 3;
    }

    unsigned int hash_len;
    if(calc_sha1(filename, &hash, &hash_len)){
        TRACE_ERROR("calc_sha1 error\n");
        return 3;
    }

    /* load signature from file */
    BIO *in = NULL;
    in = BIO_new(BIO_s_file());
    if(in == NULL) {
        TRACE_ERROR("BIO_new error\n");
        return 1;
    }

    if(!BIO_read_filename(in, (char*)sig_file)) {
        TRACE_ERROR("BIO_read_filename error for file: %s\n", sig_file);
        return 2;
    }

    unsigned char *signature = (unsigned char *)malloc(MAX_SIG_LEN);
    int signature_len = BIO_read(in, signature, MAX_SIG_LEN);

    if(signature_len <= 0) {
        TRACE_ERROR("problem reading signature from file\n");
        return 1;
    }

    DSA_SIG *dsa_sig = DSA_SIG_new();
    if(dsa_sig == NULL) {
        TRACE_ERROR("DSA_SIG_new fail\n");
        return 1;
    }

    if(NULL == d2i_DSA_SIG(&dsa_sig, (const unsigned char **) &signature,
            signature_len)){
        TRACE_ERROR("Decoding signature problem\n");
        return 1;
    }

    /* compute k
     * k == H(m)*inv(s) + x*r*inv(s)  (mod q)
     * k == H(m)*w + x*r*w (mod q)
     */

    /* compute w == inv(s) (mod q) */
    BN_init(w);
    if ((BN_mod_inverse(w, dsa_sig->s, dsa->q, ctx)) == NULL){
        TRACE_ERROR("BN_mod_inverse error\n");
        return 1;
    }

    /* H(m) */
    BN_init(m);
    if (BN_bin2bn(hash, hash_len, m) == NULL) {
        TRACE_ERROR("BN_bin2bn error\n");
        return 1;
    }

    /* sequentially compute k */
    /* k = xr */
    if (!BN_mod_mul(k, dsa->priv_key, dsa_sig->r, dsa->q, ctx)) {
        TRACE_ERROR("BN_mod_mul error\n");
        return 1;
    }

    /* k = m + xr */
    if (!BN_add(k, k, m)) {
        TRACE_ERROR("BN_add error\n");
        return 1;
    }

    if (!BN_mod(k, k, dsa->q, ctx)) {
        TRACE_ERROR("BN_mod error\n");
        return 1;
    }

    /* k = w(m + xr) */
    if (!BN_mod_mul(k, k, w, dsa->q, ctx)) {
        TRACE_ERROR("BN_mod_mul error\n");
        return 1;
    }

    /* return k into string */
    // out buffer to be freed by the caller
    *msg = (unsigned char *)malloc(21);
    memset(*msg, 0, 21);

    BN_bn2bin(k, *msg);

    BN_free(k);
    BN_free(m);
    BN_free(w);

    return 0;
}
